import copy
import torch
from torch.utils import data
import os
import numpy as np
import logging
from sklearn import metrics
# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)

'''
约定上，客户端必须继承该类，避免使用强制接口。并实现以下方法；
1. __init__()初始化自己的客户端
2. set_parameters()加载本地的特殊参数
3. train() 用来完成特异化的训练
4. test() 用来完成特异化的测试

train/test 因为没有附带调用关系（由父类的方法调用，只在子类内部调用）所以可以重新定义，不需要复写，用来区分。
如果有父类子类的名字不同，应该防止出现错误。不小心使用了父类中的该方法。
其他的函数。可以直接写在基类里作为通用的方法。
'''

class FedClientBase(object):
    """训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。\n
    这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。

    1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
    2. 训练的模型、参数在客户端配置。
    3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。
    接收参数，完成训练。主要是一次训练的参数。

    每一个算法都应该提供梯度的返回值。即在客户端计算出梯度。其实在服务器计算也无可厚非。
    Args:
        object ([type]): [description]
    """    


    def __init__(self,datapair,model):
        """创建一个客户端对象。初始化客户端的本地参数。作为基类。初始化所有的基类参数，与模型无关的参数。
        1. device 设备，检查GPU是否可用
        2. user_id 随机生成的客户端标识
        3. model_path 用来存储和加载模型的路径
        4. data_pair x,y本地数据集，ndarray格式
        5. model 模型
        6. optimizer 优化器

        Args:
            datapair (tuple(ndarray,ndarray)): input,output包括输入输出数据
            model (torch.model): 与客户端一致的模型
        """        
        # 客户端本地参数-------------------------------------------------------------------------------
        # 当前可以用来训练的设备
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        # 当前用户的一个随机id（八位整数，分布式情况下可以使用IP地址：port端口号作为id）
        self.user_id = torch.randint(10000000,99999999,(1,)).item()
        # 用来持久化当前的模型的路径
        self.model_path = "models/"

        # 训练设备、训练数据、训练模型、损失函数、优化器----------------------------------------------------------
        self.datapair = datapair
        self.model = copy.deepcopy(model).to(self.device)
        self.loss_func = torch.nn.CrossEntropyLoss()
        self.optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01)
        # 需要进行参数聚合的Adam优化器
        # self.optimizer = optim.Adam(self.model.parameters(),lr=0.001)
        # 不需要进行参数聚合的SGD优化器
        # self.optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01)
        # 带有正则化项的Adam优化器
        # self.optimizer = optim.Adam(self.model.parameters(),lr=0.001,weight_decay=0.000001)
        # 带有正则化项的SGD优化器
        # self.optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01,weight_decay=0.000001)
        logger.info("FedClientBase init:device:{}\tuser_id:{}".format(self.device,self.user_id))

    def set_parameters(self,client_params):
        """从服务器加载训练的参数。\n
        1. epochs 数据集循环的次数
        2. batch_size 每个batch的大小
        3. inner_lr 客户端优化器的步长
        4. outer_lr 服务器端优化器的步长
        5. update_steps 本地更新的步骤。如果定义了这个变量，则batch_size无效。batch_size=len(dataset)/update_steps
        Args:
            client_params (dict): 客户端进行训练的相关参数
        """        

        # 训练参数---------------------------------------------------------------------------------------
        # 基本训练参数
        if("epochs" in client_params):
            self.epochs = client_params["epochs"]
        else:
            self.epochs =1
        
        if("batch_size" in client_params):
            self.batch_size = client_params["batch_size"]
        else:
            self.batch_size=128
        if("inner_lr" in client_params):
            self.inner_lr=client_params["inner_lr"]
        else:
            self.inner_lr = 0.01
        if("outer_lr" in client_params):
            self.outer_lr=client_params["outer_lr"]
        else:
            self.outer_lr = 0.8

        # 根据innerlr修改optimizer
        optimizer_state_dict = self.optimizer.state_dict()
        optimizer_state_dict['param_groups'][0]['lr']=self.inner_lr
        self.optimizer.load_state_dict(optimizer_state_dict)

        # 并且本地update_steps重新规定batchsize
        if "update_steps" in client_params :
            self.update_steps=client_params["update_steps"]
        else:
            self.update_steps = None
        if self.update_steps != None:
            self.batch_size =int (len(self.datapair[1])/self.update_steps)

        logger.info("FedClientBase set_parameters:epochs={}\t batch_size={}\t update_steps={}".format(self.epochs,self.batch_size,self.update_steps))


    def process_data(self):
        """处理数据self.datapair,生成self.traindl\n
        什么样的函数应该用来提供返回值，什么样的函数，应该动态添加类变量。果然是个很难理解的东西。\n
        我觉得。只要写在类的内部的函数，应该都可以直接修改类的状态。而不是通过返回值让其他的函数进行修改。\n
        甚至可以不用传递入参。如果需要传递入参出参数。而不是直接访问类内部的变量或者修改类内部的变量，那么直接定义工具util函数即可。\n
        所以，这里完全可以不用入参传。类内函数，与类内的变量就应该高度耦合。通过出参数如参数解耦的方法，应该是类外的函数使用。\n

        Returns:
            train_dl: 封装好的本地数据
        """        
        x,y = self.datapair
        assert len(x) == len(y)
        # 将数据转换成tensor,为了加快速度，通过共享内存的方式
        x_train_tensor = torch.FloatTensor(x).to(self.device)
        y_train_tensor = torch.from_numpy(y).to(self.device)
        # print(x_train_tensor.dtype,y_train_tensor.dtype)
        train_ds = data.TensorDataset(x_train_tensor,y_train_tensor)
        self.train_dl = data.DataLoader(train_ds,batch_size=self.batch_size,shuffle=True)

        # 每一个round包含多少个iteration
        self.iteration  = self.epochs*(len(self.train_dl))
        logger.info("FedClientBase process_data: iteration:{}\t".format(self.iteration))
        return self.train_dl

    
    def process_data_support_and_query(self):
        """注意，这里将support&&query set都设置为一个batch，相当于one batch梯度下降。
        """        
        x,y = self.datapair
        assert len(x) == len(y)
        sq_point  = int(len(y)*self.support_ratio)
        # 将数据转换成tensor,为了加快速度，通过共享内存的方式
        # x_support_tensor = torch.FloatTensor(x[:sq_point]).to(self.device)
        x_support_tensor = torch.FloatTensor(x[:sq_point]).to(self.device)
        y_support_tensor = torch.from_numpy(y[:sq_point]).to(self.device)
        x_query_tensor = torch.FloatTensor(x[sq_point:]).to(self.device)
        y_query_tensor = torch.from_numpy(y[sq_point:]).to(self.device)
        # print(x_train_tensor.dtype,y_train_tensor.dtype)
        support_ds = data.TensorDataset(x_support_tensor,y_support_tensor)
        self.support_dl = data.DataLoader(support_ds,batch_size=self.batch_size,shuffle=True)
        query_ds  = data.TensorDataset(x_query_tensor,y_query_tensor)
        self.query_dl = data.DataLoader(query_ds,batch_size=len(query_ds),shuffle=True)
        self.iteration  = self.epochs*(len(self.support_dl))
        logger.info("FedPerClient process_data_support_and_query: iteration_support_dl:{}\t".format(self.iteration))
        return self.support_dl,self.query_dl


    def load_model_state_dict(self,model_state_dict):
        """加载model

        Args:
            model_state_dict (dict[tensor]): pytorch格式的参数
        """        
        self.model.load_state_dict(model_state_dict)


    
    def get_model_state_dict(self):
        """返回当前的模型参数字典

        Returns:
            dict[tensor]: pytorch格式的模型参数
        """        
        return self.model.state_dict()

    def load_optimizer_state_dict(self,optimizer_state_dict):
        """加载optimizer（暂时已经弃用）

        Args:
            optimizer_state_dict (dict[tensor]):pytorch格式的参数
        """        
        self.optimizer.load_state_dict(optimizer_state_dict)

    def get_optimizer_state_dict(self):
        """返回当前的优化器参数字典

        Returns:
            dict[tensor]: pytorch格式的优化器参数
        """        
        return self.optimizer.state_dict()

    def get_grad_dict(self):
        """返回当前的梯度参数列表

        Returns:
            list[tensor]:梯度字典。在以梯度为核心的优化中起到作用 
        """        
        return copy.deepcopy(self.grad_dict)
    
    def save_model(self,model_path):
        """持久化模型

        Args:
            model_path (str): 模型的存储路径
        """        
        torch.save(self.model, model_path)

    def read_model(self,model_path):
        """加载预训练的模型

        Args:
            model_path (str): 模型的存储路径
        """        
        self.model = torch.load(model_path)

    # 
    def get_params_norm(self,old_param_list,new_param_list):
        """计算两组值之间的方差。表示变化率或者说贡献率。这个是工具函数。与类没哟耦合

        Args:
            old_param_list (list[tensor]): 原有的参数
            new_param_list (list[tensor]): 最新的参数

        Returns:
            float: 梯度的二范数.即tensor向量的模长。即F范数。
        """        
        params = []

        for old_param,new_param in zip(old_param_list,new_param_list):
            params.append(torch.flatten(new_param  - old_param))

        # return torch.linalg.norm(torch.cat(params), 2)
        return float(torch.norm(torch.cat(params)))
    

    
    def get_imformation_increase(self,old_param,new_param,accuracy):
        """计算训练过程中的信息增益速度。用来衡量贡献度。待会进行计算。考虑到计算的准确率

        Args:
            old_param (list[tensor]): 原来参数
            new_param (list[tensor]): 本次参数
            accuracy (float): 准确率
        """        
        pass

    def train(self):
        """用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy

        Returns:
            [(running_loss,running_accuracy)]: 训练损失和训练的准确率
        """     

        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0
        for epoch in range(self.epochs):
            for batch_x,batch_y in self.train_dl:
                # 正向传播，计算准确度
                outputs = self.model(batch_x)
                _,predicted = torch.max(outputs,1)
                correct_sum += (predicted == batch_y).sum().item()
                total_size += len(batch_y)
                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                loss_sum +=loss.item()
                loss_times+=1
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()
                # 梯度下降
                self.optimizer.step()

        # print(self.optimizer.state_dict())
        return (loss_sum/loss_times,100*correct_sum/total_size)

    def test(self):
        """用来测试训练的结果返回准确率

        Returns:
            float: 模型在本地数据上进行测试的准确率
        """        
        self.model.eval()
        correct = 0
        total =0

        # 引入numpy主要是为了应用sklearn中的一系列算法
        y_pred = np.array([])
        y_true = np.array([])

        with torch.no_grad():
            for features,labels in self.train_dl:
                outputs = self.model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                y_pred = np.append(y_pred,predicted.numpy())
                y_true = np.append(y_true,labels.numpy())
                # total += labels.size(0)
                # correct += (predicted == labels).sum().item()

        accuracy = self.get_accuracy_score(y_true,y_pred)
        # average_precision_score = self.get_average_precision_score(y_true,y_pred)
        # average_recall_score = self.get_average_recall_score(y_true,y_pred)
        # average_f1_score =self.get_average_f1_score(y_true,y_pred)
        # hamming_loss = self.get_hamming_loss(y_true,y_pred)
        # confusion_matrix = self.get_confusion_matrix(y_true,y_pred)
        # classification_report = self.get_classification_report(y_true,y_pred)

        # precision = self.get_precision_score(y_true,y_pred)
        # recall = self.get_recall_score(y_true,y_pred)
        # 统一返回多个值。方便以后扩展，传递多个评估方案的结果。
        return ((accuracy*100).tolist(),)

    # 1 一些常见的评估指标
    # 准确率
    def get_accuracy_score(self,y_true,y_pred):
        return metrics.accuracy_score(y_true,y_pred)

    # 平均精确率
    def get_average_precision_score(self,y_true,y_pred):
        return metrics.precision_score(y_true,y_pred,average='macro')
    
    # 平均召回率
    def get_average_recall_score(self,y_true,y_pred):
        return metrics.recall_score(y_true, y_pred,average='macro')


    #  平均f1（类别不平衡的时候）
    def get_average_f1_score(self,y_true,y_score):
        return metrics.f1_score(y_true,y_score,average='macro')
    
    # 3计算相似度（有多种log_loss,hamming_loss,jaccard_similarity_score）
    def get_hamming_loss(self,y_true,y_pred):
        return metrics.hamming_loss(y_true,y_pred)
    # 4计算混淆矩阵
    def get_confusion_matrix(self,y_true,y_pred):
        return metrics.confusion_matrix(y_true, y_pred)
    
    # 5得到分类报告
    def get_classification_report(self,y_true,y_pred):
        return metrics.classification_report(y_true,y_pred)

    # 6 f1score二分类的时候使用

    # 精确率
    def get_precision_score(self,y_true,y_pred):
        return metrics.precision_score(y_true,y_pred)
    
    # 召回率
    def get_recall_score(self,y_true,y_pred):
        return metrics.recall_score(y_true, y_pred,)
    
    def get_f1_score(y_true,y_pred):
        return metrics.f1_score(y_true, y_pred)
    
    def get_roc_auc_score(y_true,y_pred):
        return metrics.roc_auc_score(y_true,y_pred)